{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "def points_collect(event,x,y,flags,param):\n",
    "    dic_points = param\n",
    "    if event == cv2.EVENT_LBUTTONDBLCLK:\n",
    "        if len(dic_points['p2'])>0:\n",
    "            dic_points['p2']=()\n",
    "            dic_points['p1']=(x,y)\n",
    "        elif len(dic_points['p1'])==0:\n",
    "            dic_points['p1']=(x,y)\n",
    "        elif len(dic_points['p1'])>0:\n",
    "            dic_points['p2']=(x,y)\n",
    "            \n",
    "    if event == cv2.EVENT_MOUSEMOVE:\n",
    "        dic_points['p_move']=(x,y)\n",
    "        \n",
    "        \n",
    "def drawline(img,dic_points):\n",
    "    color = (0,255,255)\n",
    "    if len(dic_points['p2'])>0:\n",
    "        cv2.circle(img,dic_points['p1'],2,color,cv2.FILLED)\n",
    "        cv2.circle(img,dic_points['p2'],2,color,cv2.FILLED)\n",
    "        cv2.line(img,dic_points['p1'],dic_points['p2'],color,1)\n",
    "    elif len(dic_points['p1'])>0 and len(dic_points['p_move'])>0:\n",
    "        cv2.circle(img,dic_points['p1'],2,color,cv2.FILLED)\n",
    "        cv2.line(img,dic_points['p1'],dic_points['p_move'],color,1)\n",
    "\n",
    "# 图像叠加\n",
    "def img_overlayer(img,img_fg,pos_fg,bk_fg):\n",
    "    \n",
    "    #把前景图变换为灰度\n",
    "    fg_gray = cv2.cvtColor(img_fg,cv2.COLOR_BGR2GRAY)\n",
    "    h_gf,w_fg = np.shape(fg_gray)\n",
    "    \n",
    "    \n",
    "    # 获取前景图的mask 有图部分为 1 背景部分为 0 \n",
    "    if bk_fg == 255:\n",
    "        mask_fg = fg_gray<250\n",
    "    elif bk_fg == 0:\n",
    "        mask_fg = fg_gray>5\n",
    "    \n",
    "    mask_fg = mask_fg[:,:,np.newaxis]\n",
    "    not_mask_fg = ~mask_fg\n",
    "    \n",
    "    # 截取背景图\n",
    "    bk = img[pos_fg[1]:pos_fg[1]+h_gf,pos_fg[0]:pos_fg[0]+w_fg]\n",
    "    \n",
    "    img_overlayer = bk*not_mask_fg + img_fg*mask_fg\n",
    "    img[pos_fg[1]:pos_fg[1]+h_gf,pos_fg[0]:pos_fg[0]+w_fg] = img_overlayer\n",
    "    return img\n",
    "    \n",
    "def add_glass(img,img_glass,dic_points,bk_fg=255):\n",
    "    \n",
    "    if len(dic_points['p2'])==0:\n",
    "        return img\n",
    "    \n",
    "    # 获取灭火器图像大小\n",
    "    w_glass = np.shape(img_glass)[1]\n",
    "    h_glass = np.shape(img_glass)[0]\n",
    "    \n",
    "    # 计算缩放尺度\n",
    "    scale = np.abs(dic_points['p1'][0]-dic_points['p2'][0])/w_glass\n",
    "    \n",
    "    # 图像缩放\n",
    "    resize_glass = cv2.resize(img_glass,(int(w_glass*scale),int(h_glass*scale)))\n",
    "\n",
    "    # 计算灭火器图像的起始位置(左上坐标)\n",
    "    pos_glass = (dic_points['p1'][0],dic_points['p1'][1]-int(h_glass*scale/2.0))\n",
    "    \n",
    "    \n",
    "    img_out = img_overlayer(img,resize_glass,pos_glass,bk_fg)\n",
    "    \n",
    "    return img_out\n",
    "    \n",
    "  \n",
    "if __name__==\"__main__\":\n",
    "    \n",
    "    # 记录坐标的字典\n",
    "    dic_points = {}\n",
    "    dic_points[\"p1\"]=()\n",
    "    dic_points[\"p2\"]=()\n",
    "    dic_points[\"p_move\"]=()\n",
    "    \n",
    "    # 设置回调函数\n",
    "    cv2.namedWindow('image')\n",
    "    cv2.setMouseCallback('image', points_collect,param=dic_points)\n",
    "    \n",
    "    # 不加灭火器\n",
    "    flag_add_glass = 0\n",
    "    \n",
    "    # 读取灭火器图像\n",
    "    img_glass = cv2.imread(\"fire.bmp\")\n",
    "    \n",
    "    while True:\n",
    "        img = cv2.imread(\"home.jpg\")  #加载图片 \n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "        if flag_add_glass ==0:\n",
    "            drawline(img,dic_points)\n",
    "        \n",
    "        if flag_add_glass ==1:\n",
    "            img = add_glass(img,img_glass,dic_points)\n",
    "        \n",
    "        key = cv2.waitKey(1) & 0xFF\n",
    "        \n",
    "        if key == ord('q'):\n",
    "            break\n",
    "        if key == ord('a'):\n",
    "           flag_add_glass = 1\n",
    "           key =0\n",
    "        if key == ord('r'):\n",
    "           flag_add_glass = 0\n",
    "           key =0\n",
    "        \n",
    "        cv2.imshow(\"image\",img)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:Anaconda3]",
   "language": "python",
   "name": "conda-env-Anaconda3-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
